import pickle as pkl
import pandas as pd


def getAllNewsTypes(data):
    types = set([])
    for tps in data.iloc[:, 2]:
        types.update(tps.split('|'))
    return list(types)


news_header = ['新闻id', '新闻名', '新闻类型']
news_data = pd.read_csv('./data/news_data/news.dat', delimiter='::', header=None, names=news_header, engine='python')

news_types = getAllNewsTypes(news_data)


user_header = ['用户id', '性别', '年龄', '职业', '邮编']
user_data = pd.read_csv('./data/news_data/users.dat', delimiter='::', header=None, names=user_header, engine='python')

rating_header = ['用户id', '新闻id', '评分', '时间戳']
rating_data = pd.read_csv('./data/news_data/ratings.dat', delimiter='::', header=None, names=rating_header, engine='python')


def getNewsDetailDict(data):
    '''获取news详情字典.

    格式：
    {
        news_id: {      # 新闻id, int
           news_title,  # 新闻名, string
           news_types   # 新闻类型,list
        }
    }

    '''
    ret = {}
    for index, row in data.iterrows():
        ret[row['新闻id']] = {
            "news_title": row['新闻名'],
            "news_types": row['新闻类型'].split('|')
        }

    return ret


def getUserDetailDict(data):
    '''获取用户详情字典.

    格式:
    {
        user_id: {
            gender,   # 性别, string
            age,      # 年龄段, int
            job,      # 职业, int
        }
    }
    '''
    ret = {}
    for index, row in data.iterrows():
        ret[row['用户id']] = {
            "gender": row['性别'],
            "age": row['年龄'],
            "job": row['职业']
        }

    return ret


def getNewsTitleDict(data):
    '''获取新闻名词典和最长的新闻名长度.

    将所有的新闻的新闻名看作一个词袋（word bag），例如： "love story"长度为2

    格式:
    {
        term: index
    }
    '''
    index = 1
    max_len = 0
    ret = {}
    for _, row in data.iterrows():
        title = row['新闻名']
        terms = title.split(' ')
        max_len = max(max_len, len(terms))
        for term in terms:
            if term not in ret:
                ret[term] = index
                index += 1

    return ret, max_len


n_user = max(user_data['用户id'])  # 用户数目（最大用户id）
n_news = max(news_data['新闻id'])  # 新闻数目（最大新闻id）

news_detail_dict = getNewsDetailDict(news_data)  # 新闻详情
user_detail_dict = getUserDetailDict(user_data)  # 用户详情

user_age_sep = [1, 18, 25, 35, 45, 50, 56]  # 用户年龄段
user_job_sep = [i for i in range(21)]  # 用户工作类别
user_gender_sep = ['F', 'M']  # 用户性别类别

news_title_dict, title_max_len = getNewsTitleDict(news_data)  # 新闻名字典, 新闻名最大长度


def padding_list(lis, padding_size):
    'padding数据...'
    for _ in range(padding_size - len(lis)):
        lis.append(0)
    return lis


dataset_header = ['user_id', 'news_id', 'user_gender', 'user_age', 'user_job', 'news_title', 'news_types', 'rank']

user_id_series = rating_data['用户id']
news_id_series = rating_data['新闻id']
user_gender_series = user_id_series.map(lambda user_id: user_gender_sep.index(user_detail_dict[user_id]['gender']))
user_age_series = user_id_series.map(lambda user_id: user_age_sep.index(user_detail_dict[user_id]['age']))
user_job_series = user_id_series.map(lambda user_id: user_detail_dict[user_id]['job'])
news_title_series = news_id_series.map(lambda news_id: padding_list(
    [news_title_dict[term] for term in news_detail_dict[news_id]['news_title'].split(' ')],
    padding_size=title_max_len))
news_types_series = news_id_series.map(
    lambda news_id: padding_list([news_types.index(tp) + 1 for tp in news_detail_dict[news_id]['news_types']],
                                  padding_size=len(news_types)))
rank_series = rating_data['评分']

# 创建数据集
dataset = pd.concat(
    [
        user_id_series.rename('user_id'),
        news_id_series.rename('news_id'),
        user_gender_series.rename('user_gender'),
        user_age_series.rename('user_age'),
        user_job_series.rename('user_job'),
        news_title_series.rename('news_title'),
        news_types_series.rename('news_types'),
        rank_series.rename('rank')
    ],
    axis=1
)

print(dataset.head())

# 保存数据到本地
dataset.to_pickle('./data/data.pkl')
# pkl.dump(dataset, open('./data.pkl', 'wb'))
pkl.dump(n_user, open('./data/n_user.pkl', 'wb'))
pkl.dump(n_news, open('./data/n_news.pkl', 'wb'))
pkl.dump(news_title_dict, open('./data/news_title_dict.pkl', 'wb'))
